import numpy as np
from mindquantum import *
import matplotlib.pyplot as plt
from tqdm import trange

# give an alternative way to efficiently calculate qfi elements on classical simulator.
#partial_index is to index different parameters in the circuit, ranging from 0 to n_params.
def dqs(circ, partial_index_1, params):
    sim = Simulator('projectq', circ.n_qubits)
    params[partial_index_1] += np.pi
    sim.apply_circuit(circ, params)
    d1_qs = 0.5*sim.get_qs()
    params[partial_index_1] -= np.pi
    return d1_qs

def qs(circ, params):
    sim = Simulator('projectq', circ.n_qubits)
    sim.apply_circuit(circ, params)
    qs0 = sim.get_qs()
    return qs0

# define the qfi matrix

def qfi_qsmethod(circ, params):
    n_params = len(circ.params_name)
    qfi_matrix = np.zeros([n_params, n_params])
    qs0 = qs(circ, params)
    for i in range(n_params):
        dqs1 = dqs(circ, i, params)
        for j in range(0,i+1):
            dqs0 = dqs(circ, j, params)
            qfi_matrix[j, i] = qfi_matrix[i, j] = (np.vdot(dqs0,dqs1)-(np.vdot(dqs0,qs0))*(np.vdot(qs0,dqs1))).real
    return qfi_matrix

# define the Hadamard test circuit

def circ_htest(circ, partial_param_1, partial_param_2):
    n = circ.n_qubits

    param_index_1 = circ.params_name.index(partial_param_1)
    gate_index_1 = (2*n-1)*(param_index_1//n) + param_index_1%n + n
    gate_name_1 = circ[gate_index_1].name
    obj_q_1 = circ[gate_index_1].obj_qubits

    param_index_2 = circ.params_name.index(partial_param_2)
    gate_index_2 = (2*n-1)*(param_index_2//n) + param_index_2%n + n
    gate_name_2 = circ[gate_index_2].name
    obj_q_2 = circ[gate_index_2].obj_qubits

    if gate_name_1 == 'RX':
        circ =circ[:(gate_index_1+1)] +H.on(n)+X.on(n)+ X.on(obj_q_1,n)+X.on(n) +circ[(gate_index_1+1):(gate_index_2+1)]
    elif gate_name_1 == 'RY':
        circ =circ[:(gate_index_1+1)] +H.on(n)+X.on(n)+ Y.on(obj_q_1,n)+X.on(n) +circ[(gate_index_1+1):(gate_index_2+1)]
    elif gate_name_1 == 'RZ':
        circ =circ[:(gate_index_1+1)] +H.on(n)+X.on(n)+ Z.on(obj_q_1,n)+X.on(n) +circ[(gate_index_1+1):(gate_index_2+1)]
    else:
        print('Error in partial derivative state relative to partial_param_1')
    
    if gate_name_1 == 'RX':
            circ += X.on(obj_q_2,n)
    elif gate_name_1 == 'RY':
            circ += Y.on(obj_q_2,n)
    elif gate_name_1 == 'RZ':
            circ += Y.on(obj_q_2,n)
    else:
        print('Error in partial derivative state relative to partial_param_2')
    
    return circ

def qfi_1st_term(circ, partial_param_1, partial_param_2, params):
    n = circ.n_qubits
    circ = circ_htest(circ, partial_param_1, partial_param_2)
    sim = Simulator('projectq', n+1)
    ham = Hamiltonian(QubitOperator(f'X{n}'))
    sim.apply_circuit(circ, params[:len(circ.params_name)])
    expectation = sim.get_expectation(ham).real
    return expectation/4

def qfi_2nd_term_2(circ, partial_param_1, params):
    n = circ.n_qubits

    param_index_1 = circ.params_name.index(partial_param_1)
    gate_index_1 = (2*n-1)*(param_index_1//n) + param_index_1%n + n
    gate_name_1 = circ[gate_index_1].name
    obj_q_1 = circ[gate_index_1].obj_qubits[0]

    circ = circ[:(gate_index_1)]

    if gate_name_1 == 'RX':
        ham = Hamiltonian(QubitOperator(f'X{obj_q_1}'))
    elif gate_name_1 == 'RY':
        ham = Hamiltonian(QubitOperator(f'Y{obj_q_1}'))
    elif gate_name_1 == 'RZ':
        ham = Hamiltonian(QubitOperator(f'Z{obj_q_1}'))
    else:
        print('Error in qfi_second_term1')
    sim = Simulator('projectq', n)
    sim.apply_circuit(circ, params[:len(circ.params_name)])
    expectation = sim.get_expectation(ham)

    return expectation*(-1j/2)

def qfi_2nd_term_1(circ, partial_param_1, params):
    n = circ.n_qubits

    param_index_1 = circ.params_name.index(partial_param_1)
    gate_index_1 = (2*n-1)*(param_index_1//n) + param_index_1%n + n
    gate_name_1 = circ[gate_index_1].name
    obj_q_1 = circ[gate_index_1].obj_qubits[0]

    circ = circ[:(gate_index_1)]

    if gate_name_1 == 'RX':
        ham = Hamiltonian(QubitOperator(f'X{obj_q_1}'))
    elif gate_name_1 == 'RY':
        ham = Hamiltonian(QubitOperator(f'Y{obj_q_1}'))
    elif gate_name_1 == 'RZ':
        ham = Hamiltonian(QubitOperator(f'Z{obj_q_1}'))
    else:
        print('Error in qfi_second_term2')
    sim = Simulator('projectq', n)
    sim.apply_circuit(circ, params[:len(circ.params_name)])
    expectation = sim.get_expectation(ham)

    return expectation*(1j/2)

def qfi_element(circ, partial_param_1, partial_param_2, params):
    qfi_elem = (qfi_1st_term(circ, partial_param_1, partial_param_2, params)-qfi_2nd_term_1(circ, partial_param_1, params)*qfi_2nd_term_2(circ, partial_param_2, params)).real
    if qfi_elem < 1e-10:
        return 0
    else:
        return qfi_elem

def qfi_htmethod(circ, params):
    params_list = circ.params_name
    n_params = len(params_list)
    qfi_matrix = np.zeros([n_params, n_params])
    for i in range(n_params):
        for j in range(0,i+1):
            qfi_matrix[j, i] = qfi_matrix[i, j] = qfi_element(circ, params_list[j], params_list[i], params)
            if qfi_matrix[j, i]<1e-11:
                return 0
            else:
                return qfi_matrix

def nat_grad(ham, circ, params):
    sim = Simulator('projectq', circ.n_qubits)
    exp = sim.get_expectation_with_grad(ham, circ)
    grad = exp(params)[1][0,0]
    g_pinv = np.linalg.pinv(qfi_qsmethod(circ, params))
    nat_grad = np.matmul(g_pinv, grad)
    return nat_grad.real

#************

def RP(rotation_angle):
    a = np.random.randint(0, 3)
    if a == 0:
        return RX(rotation_angle)
    elif a == 1:
        return RY(rotation_angle)
    elif a == 2:
        return RZ(rotation_angle)
    else: 
        print("error in random Pauli gates")

def bpansatz(n,p):
    init_state_circ = UN(RY(np.pi/4),n)
    qc = Circuit() + init_state_circ
    counter = 0
    for j in range(p):
        for i in range(n):
            qc += RP(str(counter)).on(i)
            counter += 1
        for ii in range(n-1):
            qc += Z(ii,ii+1)
    return qc

# define sampling method to obtain the variance of natrial gradient.

def get_var_nat_partial_exp(circuit,  
                        hamiltonian, 
                        number_of_attempts = 10, #初始采样次数，不是实际的采样次数，程序会根据error_reques自动调整（小概率会影响精度，如果发现跑出来与预期不符，可以适当调大此项）
                        error_request = 0.3):   #要求最低的采样误差，这个在很大程度上影响代码运行时间，默认0.2会比较快，而且精度也够
    avg_variance = 0
    error = 1  #当前误差，初始值设置为最大值1
    iteration = 0  #标记迭代次数

    #接下里是初次采样
    for kk in trange(number_of_attempts):
        rdm = np.random.rand(len(circuit.params_name))*2*np.pi
        avg_variance += nat_grad(hamiltonian, circuit, rdm)[0]**2/(number_of_attempts)
        
    
    #这里是做审敛的条件循环，若不满足误差要求，会自动增加样本容量，不断采样计算，直到结果满足要求。
    while error > error_request:
        iteration += 1
        variance = 0
        for jj in trange(number_of_attempts*(2**(iteration-1))):
            rdm = np.random.rand(len(circuit.params_name))*2*np.pi
            variance += nat_grad(hamiltonian, circuit, rdm)[0]**2/(number_of_attempts*(2**(iteration-1)))
        avg_variance_i = avg_variance
        avg_variance = (variance + avg_variance)/2
        error = abs(avg_variance-avg_variance_i)/(avg_variance)
        print("已迭代次数：",iteration, "\t" "当前采样误差" ,error)
    return avg_variance

# define sampling method to obain the variance of conventional gradient.
def get_var_partial_exp(circuit,  
                        hamiltonian, 
                        number_of_attempts = 10, #初始采样次数，不是实际的采样次数，程序会根据error_reques自动调整（小概率会影响精度，如果发现跑出来与预期不符，可以适当调大此项）
                        error_request = 0.1):   #要求最低的采样误差，这个在很大程度上影响代码运行时间，默认0.2会比较快，而且精度也够
    simulator = Simulator('projectq',circuit.n_qubits)
    grad_ops = simulator.get_expectation_with_grad(hamiltonian, circuit)
    avg_variance = 0
    error = 1  #当前误差，初始值设置为最大值1
    iteration = 0  #标记迭代次数

    #接下里是初次采样
    for kk in trange(number_of_attempts):
        rdm = np.random.rand(len(circuit.params_name))*2*np.pi
        avg_variance += grad_ops(rdm)[1][0,0,0].real**2/(number_of_attempts)
    
    #这里是做审敛的条件循环，若不满足误差要求，会自动增加样本容量，不断采样计算，直到结果满足要求。
    while error > error_request:
        iteration += 1
        variance = 0
        for jj in trange(number_of_attempts*(2**(iteration-1))):
            rdm = np.random.rand(len(circuit.params_name))*2*np.pi
            variance += grad_ops(rdm)[1][0,0,0].real**2/(number_of_attempts*(2**(iteration-1)))
        avg_variance_i = avg_variance
        avg_variance = (variance + avg_variance)/2
        error = abs(avg_variance-avg_variance_i)/(avg_variance)
        print("已迭代次数：",iteration, "\t" "当前采样误差" ,error)
    return avg_variance


if __name__ == '__main__':
    n_max = 6 # maximal number of  qubits
    p = 30 # number of layers

    xxx = np.arange(2,n_max+1,2) #用于绘图时候的横坐标
    yyy = np.zeros(int(n_max/2))        #储存梯度方差的向量
    hamiltonian = Hamiltonian(QubitOperator('Z0 Z1'))
    for n in range(2, n_max+1, 2):
        circuit = bpansatz(n,p)
        print("n = ",n)
        yyy[int(n/2)-1] = get_var_partial_exp(circuit, hamiltonian)
    plt.subplot(121)
    plt.semilogy(xxx, yyy)
    plt.title('regular gradient version')
    plt.xlabel('number of qubits')
    plt.ylabel('Var(∂E)')

    xxx = np.arange(2,n_max+1,2) #用于绘图时候的横坐标
    yyy = np.zeros(int(n_max/2))        #储存梯度方差的向量
    hamiltonian = Hamiltonian(QubitOperator('Z0 Z1'))
    for n in range(2, n_max+1, 2):
        circuit = bpansatz(n,p)
        print("n = ",n)
        yyy[int(n/2)-1] = get_var_nat_partial_exp(circuit, hamiltonian)
    #在半指数图上绘制图像
    plt.subplot(122)
    plt.semilogy(xxx, yyy)
    plt.title('natrual gradient version')
    plt.xlabel('number of qubits')
    plt.suptitle('Var(∂E) vs num of qubits')
    plt.subplots_adjust(left=None, bottom=0.15, right=None, top=None, wspace=0.3, hspace=0.15)#wspace 子图横向间距， hspace 代表子图间的纵向距离，left 代表位于图像不同位置

    plt.savefig('result_var_vs_qubis.svg')
